{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pickle\n",
    "from copy import deepcopy\n",
    "import itertools\n",
    "import datetime\n",
    "import torch\n",
    "import pprint as pp\n",
    "\n",
    "import sys, os\n",
    "sys.path.append(os.path.join(os.path.dirname(\"__file__\"), '..', '..'))\n",
    "from AI_physicist.theory_learning.models import Loss_Fun_Cumu\n",
    "from AI_physicist.theory_learning.theory_model import Theory_Training, get_loss, get_best_model_idx, get_preds_valid, get_valid_prob\n",
    "from AI_physicist.theory_learning.theory_hub import Theory_Hub, load_model_dict_at_theory_hub, unification_symbolic\n",
    "from AI_physicist.theory_learning.util_theory import plot_theories, plot3D, process_loss, plot_loss_record, to_one_hot, plot_indi_domain, get_mystery, to_pixels, standardize_symbolic_expression, check_expression_matching\n",
    "from AI_physicist.theory_learning.util_theory import prepare_data_from_matrix_file, get_epochs, export_csv_with_domain_prediction, extract_target_expression, get_coeff_learned_list, load_theory_training, get_dataset_from_file, get_piecewise_dataset\n",
    "from AI_physicist.settings.filepath import theory_PATH\n",
    "from AI_physicist.settings.global_param import COLOR_LIST\n",
    "from AI_physicist.pytorch_net.util import Loss_Fun, make_dir, Loss_with_uncertainty, Early_Stopping, record_data, plot_matrices, get_args, serialize, to_string, filter_filename, to_np_array, to_Variable, get_variable_name_list, get_param_name_list\n",
    "from AI_physicist.pytorch_net.net import Conv_Autoencoder, load_model_dict, train_simple\n",
    "import pandas as pd\n",
    "pd.set_option('max_colwidth', 400)\n",
    "\n",
    "try:\n",
    "    %matplotlib inline\n",
    "    isplot = True\n",
    "except:\n",
    "    import matplotlib\n",
    "    matplotlib.use('Agg')\n",
    "    isplot = False\n",
    "is_cuda = torch.cuda.is_available()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Settings for which unification methods to turn on:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "is_symbolic_unification = True\n",
    "is_network_unification = True"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load theory hub:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Load the theory hub. Should point to the directory that the theory_hub lives.\n",
    "An example theory hub can be downloaded at https://drive.google.com/file/d/1I0P0MUjP9b_HGMacNZ0avO_0CtTtG_y0/view?usp=sharing,\n",
    "after downloading, put it under the directory dirname.\n",
    "\"\"\"\n",
    "exp_id = \"test\"\n",
    "date_time = \"10-15\"\n",
    "dirname = theory_PATH + \"/{0}_{1}/\".format(exp_id, date_time)\n",
    "filename_hub = dirname + \"file_continuous_num_4_pred_8_linear_dom_8_onehot_mse_1e-07_sim_DLs-0-3-relative_optim_adam-0.005_adam-0.001_reg_1e-08_1e-05_batch_2000_core_DLs_order_-1_lossd_None_False_infl_False_#_20000_mul_5000_MDL_both_2D_3L_id_210_7_hub.p\"\n",
    "theory_hub = load_model_dict_at_theory_hub(pickle.load(open(filename_hub, \"rb\")), is_cuda = is_cuda)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Symbolic unification:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Implementing the symbolic unification algorithm (Alg. 4) in (Wu and Tegmark, 2019):\"\"\"\n",
    "if is_symbolic_unification:\n",
    "    num_clusters = 3  # The number of clusters for clustering the theories into\n",
    "    df, exprs_unified_list = unification_symbolic(theory_hub.theory, num_clusters=num_clusters, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning network unification:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"Directly learn multiple master theories that unifies subsets of theories in the theory hub (new feature).\n",
    "Here the unification is using generalized-mean loss as introduced in (Wu and Tegmark, 2019). Instead of operating\n",
    "on data points, this generalized-mean loss is operating on theories, s.t. for each master theories, it only \n",
    "specilizes to subset of theories.\n",
    "\"\"\"\n",
    "if is_network_unification:  \n",
    "    # Settings on theory:\n",
    "    array_id = 0\n",
    "    input_size = 4\n",
    "    output_size = 2\n",
    "\n",
    "    num_master_theories = 5\n",
    "    master_model_type = \"regulated-Net\"\n",
    "    statistics_output_neurons = 6\n",
    "    master_loss_combine_mode = \"on-loss\"\n",
    "    master_loss_core = \"mse\"\n",
    "    master_loss_mode = \"harmonic\"\n",
    "    master_model_num_neurons = 8\n",
    "    master_optim_type = (\"adam\", 1e-3)\n",
    "    master_optim_type_classifier = (\"adam\", 5e-5)\n",
    "    master_reg_amp = 1e-3\n",
    "    master_reg_amp_classifier = 1e-3\n",
    "    theory_remove_threshold = 5e-3\n",
    "    target_array_id = \"7\"\n",
    "    filter_words = []\n",
    "\n",
    "    exp_id = get_args(exp_id, 1)\n",
    "    num_master_theories = get_args(num_master_theories, 2, \"int\")\n",
    "    master_model_type = get_args(master_model_type, 3)\n",
    "    statistics_output_neurons = get_args(statistics_output_neurons, 4, \"int\")\n",
    "    master_loss_combine_mode = get_args(master_loss_combine_mode, 5)\n",
    "    master_loss_core = get_args(master_loss_core, 6)\n",
    "    master_loss_mode = get_args(master_loss_mode, 7)\n",
    "    master_model_num_neurons = get_args(master_model_num_neurons, 8, \"int\")\n",
    "    master_optim_type = get_args(master_optim_type, 9, \"tuple\")\n",
    "    master_optim_type_classifier = get_args(master_optim_type_classifier, 10, \"tuple\")\n",
    "    master_reg_amp = get_args(master_reg_amp, 11, \"float\")\n",
    "    master_reg_amp_classifier = get_args(master_reg_amp_classifier, 12, \"float\")\n",
    "    theory_remove_threshold = get_args(theory_remove_threshold, 13, \"float\")\n",
    "    target_array_id = get_args(target_array_id, 14, \"int\")\n",
    "    filter_words = get_args(filter_words, 15, \"tuple\")\n",
    "    date_time = get_args(date_time, 16)\n",
    "    array_id = get_args(array_id, 17, \"int\")\n",
    "\n",
    "    # Setting up dirname and filename:\n",
    "    load_previous = True\n",
    "    filename_hub_cand = filter_filename(dirname, include = [target_array_id, \"hub.p\", *filter_words])\n",
    "    filename_hub = dirname + filename_hub_cand[0]\n",
    "    print(\"filename_hub: {0}\".format(filename_hub))\n",
    "\n",
    "    filename = filename_hub[:-6] + \".p\"\n",
    "    make_dir(filename)\n",
    "    print(\"filename: {0}\\n\".format(filename))\n",
    "\n",
    "    filename_unification = filename[:-2] + \"/unification/num_{0}_type_{1}_statistics_{2}_{3}_core_{4}_mode_{5}_neurons_{6}_optim_{7}_{8}_reg_{9}_{10}_thresh_{11}_id_{12}\".format(\n",
    "        num_master_theories, master_model_type, statistics_output_neurons, master_loss_combine_mode, master_loss_core, master_loss_mode, master_model_num_neurons,\n",
    "        to_string(master_optim_type), to_string(master_optim_type_classifier), master_reg_amp, master_reg_amp_classifier, theory_remove_threshold, array_id,\n",
    "    )\n",
    "    make_dir(filename_unification)\n",
    "    print(\"filename_unification: {0}\\n\".format(filename_unification))\n",
    "\n",
    "    # Initialize certain parameters:\n",
    "    master_reg_dict = {\"master_model\": {\"weight\": master_reg_amp, \"bias\": master_reg_amp},\n",
    "                       \"statistics_Net\": {\"weight\": master_reg_amp, \"bias\": master_reg_amp},\n",
    "                      }\n",
    "    master_reg_dict_classifier = {\"classifier\": {\"weight\": master_reg_amp_classifier, \"bias\": master_reg_amp_classifier}}\n",
    "    struct_param_regulated_Net = [\n",
    "            [master_model_num_neurons, \"Simple_Layer\", {}],\n",
    "            [master_model_num_neurons, \"Simple_Layer\", {}],\n",
    "            [output_size, \"Simple_Layer\", {\"activation\": \"linear\"}],\n",
    "    ]\n",
    "    \n",
    "    # Load theory_hub:\n",
    "    wait_time = 1\n",
    "    wait_time_exponent = 1.2\n",
    "    while True:\n",
    "        theory_hub = load_model_dict_at_theory_hub(pickle.load(open(filename_hub, \"rb\")), is_cuda = is_cuda)\n",
    "        if len(theory_hub.theory) == 0:\n",
    "            wait_time *= wait_time_exponent\n",
    "            print(\"No theory exists in the theory_hub. Wait for {0:.1f} seconds...\".format(wait_time))\n",
    "            time.sleep(wait_time)\n",
    "        else:\n",
    "            print(\"Succesfully load theory_hub {0} with non-empty theory_collections!\".format(filename_hub))\n",
    "            break\n",
    "\n",
    "    info_dict = {}\n",
    "    # Propose master_theories:\n",
    "    theory_dict = theory_hub.get_theory_tuples(input_size = input_size)\n",
    "    master_theory_dict = theory_hub.propose_master_theories(num_master_theories = num_master_theories,\n",
    "                                                            input_size = input_size,\n",
    "                                                            output_size = output_size,\n",
    "                                                            statistics_output_neurons = statistics_output_neurons,\n",
    "                                                            master_model_type = master_model_type,\n",
    "                                                            struct_param_regulated_Net = struct_param_regulated_Net,\n",
    "                                                           )\n",
    "\n",
    "    # Fit the master_theories to all the theories:\n",
    "    data_record = theory_hub.fit_master_theory(\n",
    "        master_theory_dict = master_theory_dict,\n",
    "        theory_dict = theory_dict,\n",
    "        optim_type = master_optim_type,\n",
    "        reg_dict = master_reg_dict,\n",
    "        loss_core = master_loss_core,\n",
    "        loss_mode = master_loss_mode,\n",
    "        loss_combine_mode = master_loss_combine_mode,\n",
    "        num_iter = 1000,\n",
    "        patience = 10,\n",
    "        inspect_interval = 10,\n",
    "        isplot = isplot,\n",
    "        filename = filename_unification,\n",
    "    )\n",
    "    info_dict[\"data_record_whole\"] = deepcopy(data_record)\n",
    "    info_dict[\"master_theory_whole\"] = deepcopy({name: master_theory.model_dict for name, master_theory in master_theory_dict.items()})\n",
    "    pickle.dump(info_dict, open(filename_unification + \".p\", \"wb\"))\n",
    "\n",
    "    # Assign master theory to theories:\n",
    "    group_list = theory_hub.assign_master_theories_to_theories(master_theory_dict, theory_dict)\n",
    "    print(\"=\" * 150 + \"\\nMaster_theory assignment:\")\n",
    "    for assigned_master_theory_dict, assigned_theory_dict in group_list:\n",
    "        print(\"master_theory: {0}\".format(list(assigned_master_theory_dict.keys())[0]))\n",
    "        print(\"assigned_theories: {0}\\n\".format(list(assigned_theory_dict.keys())))\n",
    "\n",
    "    # Train each master_theory individually:\n",
    "    for i, (assigned_master_theory_dict, assigned_theory_dict) in enumerate(group_list):\n",
    "        print(\"=\" * 150)\n",
    "        print(\"Fitting {0}th assigned group:\".format(i))\n",
    "        print(\"master_theory: {0}\".format(list(assigned_master_theory_dict.keys())[0]))\n",
    "        print(\"assigned_theories: {0}\\n\".format(list(assigned_theory_dict.keys())) + \"=\" * 150 + \"\\n\")\n",
    "        master_theory_name = list(assigned_master_theory_dict.keys())[0]\n",
    "        master_theory = assigned_master_theory_dict[master_theory_name]\n",
    "\n",
    "        # Train the master_model:\n",
    "        data_record = theory_hub.fit_master_theory(\n",
    "            master_theory_dict = assigned_master_theory_dict,\n",
    "            theory_dict = assigned_theory_dict,\n",
    "            optim_type = master_optim_type,\n",
    "            reg_dict = master_reg_dict,\n",
    "            loss_core = master_loss_core,\n",
    "            loss_mode = master_loss_mode,\n",
    "            loss_combine_mode = master_loss_combine_mode,\n",
    "            num_iter = 10000,\n",
    "            patience = 20,\n",
    "            inspect_interval = 50,\n",
    "            isplot = isplot,\n",
    "            filename = filename_unification,\n",
    "        )\n",
    "        info_dict[\"data_record_{0}\".format(i)] = deepcopy(data_record)\n",
    "\n",
    "        # Train the master classifier:\n",
    "        data_record_classifier = theory_hub.fit_master_classifier(\n",
    "            master_theory = master_theory,\n",
    "            theory_dict = assigned_theory_dict,\n",
    "            optim_type_classifier = master_optim_type_classifier,\n",
    "            reg_dict_classifier = master_reg_dict_classifier,\n",
    "        )\n",
    "        info_dict[\"data_record_classifier_{0}\".format(i)] = deepcopy(data_record_classifier)\n",
    "\n",
    "        # Add master_theory_tuple:\n",
    "        theory_hub = load_model_dict_at_theory_hub(pickle.load(open(filename_hub, \"rb\")), is_cuda = is_cuda)\n",
    "        theory_hub.add_master_theory(name = master_theory_name,\n",
    "                                     master_theory = assigned_master_theory_dict[master_theory_name],\n",
    "                                     theory_tuples = assigned_theory_dict,\n",
    "                                     is_replace = True,\n",
    "                                    )\n",
    "        master_theory_tuple = Master_Theory_Tuple(master_theory = assigned_master_theory_dict[master_theory_name], theory_tuples = assigned_theory_dict)\n",
    "        info_dict[\"master_theory_tuple_{0}\".format(i)] = deepcopy(master_theory_tuple.model_dict)\n",
    "\n",
    "        # Removed passed theory (whose loss with the master_theory is less than the theory_remove_threshold):\n",
    "        master_loss_fun = Master_Loss_Fun(core = master_loss_core, cumu_mode = master_loss_mode)\n",
    "        loss_matrix = master_loss_fun.get_loss_matrix(assigned_master_theory_dict, assigned_theory_dict, use_train = False)\n",
    "        passed_theory = (loss_matrix < theory_remove_threshold).data.long()[0]\n",
    "        passed_theory_names = []\n",
    "        for j in range(len(passed_theory)):\n",
    "            is_pass = passed_theory[j]\n",
    "            if is_pass == 1:\n",
    "                passed_theory_names.append(list(assigned_theory_dict.keys())[j])\n",
    "    #     if experiment_mode != \"on-unification\":\n",
    "    #         popped_theories = theory_hub.remove_theories(passed_theory_names)\n",
    "    #         pickle.dump(theory_hub.model_dict, open(filename_hub, \"wb\"))\n",
    "        info_dict[\"popped_theories\"] = passed_theory_names\n",
    "        pickle.dump(info_dict, open(filename_unification + \".p\", \"wb\"))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
